/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2022. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/**
 * Description: Service cpp generator.
 */
#include "datasystem/common/rpc/plugin_generator/zmq_rpc_generator.h"

namespace datasystem {

void ZmqRpcGenerator::CreateServiceCpp(const google::protobuf::FileDescriptor &file,
                                       compiler::GeneratorContext *generatorCtx) const
{
    std::unique_ptr<io::ZeroCopyOutputStream> outputFile(generatorCtx->Open(fileName + ".service.rpc.pb.cc"));
    io::Printer printer(outputFile.get(), '$');

    GenerateServiceCppPrologue(printer, file);
    printer.PrintRaw(namespaceBegin);

    for (auto i = 0; i < file.service_count(); ++i) {
        auto *svc = file.service(i);
        const std::string &svcName = svc->name();
        GenerateInitMethodMapDef(printer, *svc, PREFIX, svcName);
        ImplementZmqCallMethodDef(printer, *svc, PREFIX, svcName);
    }

    printer.PrintRaw(namespaceEnd);
}

void ZmqRpcGenerator::GenerateServiceCppPrologue(io::Printer &printer,
                                                 const google::protobuf::FileDescriptor &file) const
{
    std::map<std::string, std::string> vars;
    vars["full_file_name"] = file.name();
    vars["file_name"] = fileName;

    std::string impl =
        "// Generated by the protocol buffer compiler.  DO NOT EDIT!\n"
        "// source: $full_file_name$\n"
        "#include \"$file_name$.service.rpc.pb.h\"\n"
        "#include \"datasystem/common/rdma/urma_manager_wrapper.h\"\n";
    printer.Print(vars, impl.c_str());
}

void ZmqRpcGenerator::ImplementCallMethodNoStream(io::Printer &printer,
                                                  const google::protobuf::MethodDescriptor &method, int methodIndex,
                                                  const std::string &indent)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["methodIndex"] = std::to_string(methodIndex);
    vars["methodName"] = method.name();
    vars["inputTypeName"] = method.input_type()->name();
    vars["outputTypeName"] = method.output_type()->name();
    vars["optSendPayload1"] = HasPayloadSendOption(method) ? ", std::move(payload)" : "";
    vars["optRecvPayload1"] = HasPayloadRecvOption(method) ? ", outPayload" : "";
    std::string impl =
        "            auto &methodObj = methodMap_.find($methodIndex$)->second;\n"
        "            $inputTypeName$ rq;\n"
        "            $outputTypeName$ reply;\n"
        "            auto serverApi =\n"
        "                std::make_unique<::datasystem::ServerUnaryWriterReaderImpl<$outputTypeName$, "
        "$inputTypeName$>>(\n"
        "                    sock, meta, std::move(inMsg), methodObj->HasPayloadSendOption(),\n"
        "                    methodObj->HasPayloadRecvOption());\n"
        "            rc = serverApi->Read(rq);\n"
        "            if (rc.IsError()) { break; }\n";
    if (HasPayloadSendOption(method)) {
        impl +=
            "            std::vector<::datasystem::RpcMessage> payload;\n"
            "            serverApi->ReceivePayload(payload);\n"
            "            if (rc.IsError()) { break; }\n";
    }
    if (HasPayloadRecvOption(method)) {
        impl += "            std::vector<::datasystem::RpcMessage> outPayload;\n";
    }
    impl +=
        "            rc = $methodName$(rq, reply$optSendPayload1$$optRecvPayload1$);\n"
        "            if (rc.IsError()) { rc = serverApi->SendStatus(rc); break; }\n"
        "            rc = serverApi->Write(reply);\n"
        "            if (rc.IsError()) { break; }\n";
    if (HasPayloadRecvOption(method)) {
        impl +=
            "            rc = serverApi->SendPayload(outPayload);\n"
            "            if (rc.IsError()) { break; }\n";
    }
    printer.Print(vars, impl.c_str());
}

void ZmqRpcGenerator::ImplementCallMethodClientStream(io::Printer &printer,
                                                      const google::protobuf::MethodDescriptor &method, int methodIndex,
                                                      const std::string &indent)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["methodIndex"] = std::to_string(methodIndex);
    vars["methodName"] = method.name();
    vars["inputTypeName"] = method.input_type()->name();
    vars["outputTypeName"] = method.output_type()->name();
    vars["optRecvPayload1"] = HasPayloadRecvOption(method) ? ", out" : "";
    std::string impl =
        "            auto &methodObj = methodMap_.find($methodIndex$)->second;\n"
        "            auto pimpl = std::make_unique<::datasystem::ServerReaderImpl<$inputTypeName$>>(\n"
        "                sock, meta, methodObj->HasPayloadSendOption(), methodObj->HasPayloadRecvOption(),\n"
        "                seqNo);\n"
        "            auto reader = std::make_shared<::datasystem::ServerReader<$inputTypeName$>>(std::move(pimpl));\n"
        "            $outputTypeName$ reply;\n";
    if (HasPayloadRecvOption(method)) {
        impl += "            std::vector<::datasystem::RpcMessage> out;\n";
    }
    impl +=
        "            rc = $methodName$(reader, reply$optRecvPayload1$);\n"
        "            if (rc.IsError()) { rc = reader->SendStatus(rc); break; }\n"
        "            rc = reader->WritePb(reply);\n"
        "            if (rc.IsError()) { break; }\n";
    if (HasPayloadRecvOption(method)) {
        impl +=
            "            rc = reader->SendPayload(out);\n"
            "            if (rc.IsError()) { break; }\n";
    }
    printer.Print(vars, impl.c_str());
}

void ZmqRpcGenerator::ImplementCallMethodServerStream(io::Printer &printer,
                                                      const google::protobuf::MethodDescriptor &method, int methodIndex,
                                                      const std::string &indent)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["methodIndex"] = std::to_string(methodIndex);
    vars["methodName"] = method.name();
    vars["inputTypeName"] = method.input_type()->name();
    vars["outputTypeName"] = method.output_type()->name();
    vars["optSendPayload1"] = HasPayloadSendOption(method) ? ", std::move(pl)" : "";
    std::string impl =
        "            auto &methodObj = methodMap_.find($methodIndex$)->second;\n"
        "            $inputTypeName$ rq;\n";
    if (HasPayloadSendOption(method)) {
        impl += "            std::vector<::datasystem::RpcMessage> pl;\n";
    }
    impl +=
        "            auto pimpl = std::make_unique<::datasystem::ServerWriterImpl<$outputTypeName$>>(\n"
        "                sock, meta, methodObj->HasPayloadSendOption(), methodObj->HasPayloadRecvOption(),\n"
        "                seqNo);\n"
        "            auto writer = std::make_shared<::datasystem::ServerWriter<$outputTypeName$>>(std::move(pimpl));\n"
        "            rc = writer->ReadPb(rq);\n"
        "            if (rc.IsError()) { break; }\n";
    if (HasPayloadSendOption(method)) {
        impl +=
            "            rc = writer->ReceivePayload(pl);\n"
            "            if (rc.IsError()) { break; }\n";
    }
    impl +=
        "            rc = $methodName$(writer, rq$optSendPayload1$);\n"
        "            if (rc.IsError()) { rc = writer->SendStatus(rc); break; }\n";
    printer.Print(vars, impl.c_str());
}

void ZmqRpcGenerator::ImplementCallMethodStream(io::Printer &printer, const google::protobuf::MethodDescriptor &method,
                                                int methodIndex, const std::string &indent)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["methodIndex"] = std::to_string(methodIndex);
    vars["methodName"] = method.name();
    vars["inputTypeName"] = method.input_type()->name();
    vars["outputTypeName"] = method.output_type()->name();
    std::string impl =
        "            auto &methodObj = methodMap_.find($methodIndex$)->second;\n"
        "            auto pimpl =\n"
        "                std::make_unique<::datasystem::ServerWriterReaderImpl<$outputTypeName$, $inputTypeName$>>(\n"
        "                    sock, meta, methodObj->HasPayloadSendOption(), methodObj->HasPayloadRecvOption(),\n"
        "                    seqNo);\n"
        "            auto stream =\n"
        "                std::make_shared<::datasystem::ServerWriterReader<$outputTypeName$, "
        "$inputTypeName$>>(std::move(pimpl));\n"
        "            rc = $methodName$(stream);\n"
        "            if (rc.IsError()) { rc = stream->SendStatus(rc); break; }\n";
    printer.Print(vars, impl.c_str());
}

void ZmqRpcGenerator::ImplementCallMethodUnarySocket(io::Printer &printer,
                                                     const google::protobuf::MethodDescriptor &method, int methodIndex,
                                                     const std::string &indent)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["methodIndex"] = std::to_string(methodIndex);
    vars["methodName"] = method.name();
    vars["inputTypeName"] = method.input_type()->name();
    vars["outputTypeName"] = method.output_type()->name();
    std::string impl =
        "            auto &methodObj = methodMap_.find($methodIndex$)->second;\n"
        "            auto pimpl =\n"
        "                std::make_unique<::datasystem::ServerUnaryWriterReaderImpl<$outputTypeName$, "
        "$inputTypeName$>>(\n"
        "                    sock, meta, std::move(inMsg), methodObj->HasPayloadSendOption(),\n"
        "                    methodObj->HasPayloadRecvOption());\n"
        "            auto serverApi = std::make_shared<::datasystem::ServerUnaryWriterReader<$outputTypeName$, "
        "$inputTypeName$>>(std::move(pimpl));\n"
        "            rc = $methodName$(serverApi);\n"
        "            if (rc.IsError()) { rc = serverApi->SendStatus(rc); break; }\n";
    printer.Print(vars, impl.c_str());
}

void ZmqRpcGenerator::ImplementZmqCallMethodDef(io::Printer &printer, const google::protobuf::ServiceDescriptor &svc,
                                                const std::string &indent, const std::string &svcName)
{
    const std::string &level1Indent = indent;
    const std::string level2Indent = level1Indent + indent;
    const std::string level3Indent = level2Indent + indent;
    std::map<std::string, std::string> vars;
    vars["svcName"] = svcName;
    std::string startFunction =
        "::datasystem::Status $svcName$::CallMethod(std::shared_ptr<::datasystem::ZmqServerMsgQueRef> sock,\n"
        "                                           ::datasystem::MetaPb meta,\n"
        "                                           std::deque<::datasystem::ZmqMessage> &&inMsg, int64_t seqNo) {\n"
        "    datasystem::Status rc;\n"
        "    (void)seqNo;\n"
        "    switch(meta.method_index()) {\n";
    printer.Print(vars, startFunction.c_str());
    for (auto j = 0; j < svc.method_count(); ++j) {
        if (svc.method(j) == nullptr) {
            continue;
        }
        auto &method = *(svc.method(j));
        vars["methodName"] = method.name();
        vars["methodIndex"] = std::to_string(j);
        printer.Print(vars, "        case $methodIndex$: { // $methodName$\n");
        if (!method.client_streaming() && !method.server_streaming()) {
            if (UnarySocketNeeded(method)) {
                ImplementCallMethodUnarySocket(printer, method, j, level3Indent);
            } else {
                ImplementCallMethodNoStream(printer, method, j, level3Indent);
            }
        } else if (method.client_streaming() && !method.server_streaming()) {
            ImplementCallMethodClientStream(printer, method, j, level3Indent);
        } else if (!method.client_streaming() && method.server_streaming()) {
            ImplementCallMethodServerStream(printer, method, j, level3Indent);
        } else {
            ImplementCallMethodStream(printer, method, j, level3Indent);
        }
        printer.PrintRaw(
            "            break;\n"
            "        }  // case\n");
    }
    std::string endFunction =
        "        default: {\n"
        "            rc = datasystem::Status(datasystem::StatusCode::K_UNKNOWN_ERROR, __LINE__, __FILE__,\n"
        "                                    \"Unknown method\");\n"
        "            break;\n"
        "        }\n"
        "    } // switch\n"
        "    return rc;\n"
        "}\n";
    printer.PrintRaw(endFunction);
}
}  // namespace datasystem
